import numpy as np
import gurobipy as gp
from gurobipy import GRB

import verify_piecewise_utils
from NN import cdinn2
from prob import *

# case_name = "zhrC11"

margin1 = 0.01
margin2 = 0.01

# https://support.gurobi.com/hc/en-us/articles/4414392016529-How-do-I-model-conditional-statements-in-Gurobi
def relu_der(m, x, useVars):
    eps = 0.0001
    M = 10000 + eps

    if useVars:
        relu_der_x = m.addMVar(x.shape, vtype=GRB.CONTINUOUS)
        for i in range(x.shape[0]):
            b = m.addVar(vtype=GRB.BINARY)
            m.addConstr(x[i] >= eps - M * (1 - b))
            m.addConstr(x[i] <= M * b)
            m.addConstr((b == 1) >> (relu_der_x[i] == 1))
            m.addConstr((b == 0) >> (relu_der_x[i] == 0))
        m.update()
    else:
        relu_der_x = []
        for i in range(x.shape[0]):
            if x[i] > 0:
                relu_der_x.append(1)
            else:
                relu_der_x.append(0)
        relu_der_x = np.array(relu_der_x)
    return relu_der_x


def barrier_f(m, x0, wb_list, wb_length, useVars):
    x_layer = x0
    if useVars:
        jacobian = np.eye(x0.shape[0])
    else:
        jacobian = np.eye(len(x0))
    for i in range(wb_length):
        W = wb_list[i * 2]
        b = wb_list[i * 2 + 1]
        if useVars:
            y_layer = m.addMVar((W.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
            E = np.identity(y_layer.shape[0])
            expr = W @ x_layer + b - E @ y_layer
            m.addConstr(expr == 0)
        else:
            # print("W", W.shape)
            # x_layer = np.array(x_layer)
            # print("x_layer", x_layer.shape)
            # print("b:", b.shape)
            y_layer = W @ x_layer + b
            # y_layer = np.array(y_layer)
            # print("y_layer:", y_layer.shape)
        jacobian = W @ jacobian

        if i != wb_length - 1:
            if useVars:
                z_layer = m.addMVar((y_layer.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
                # z = activate(y)
                for j in range(y_layer.shape[0]):
                    m.addConstr(z_layer[j] == gp.max_(y_layer[j], 0.0))
            else:
                z_layer = y_layer
                for j in range(y_layer.shape[0]):
                    z_layer[j] = max(y_layer[j], 0.0)
            x_layer = z_layer

            array_relu_der = []
            relu_der_f = relu_der(m, y_layer, useVars)
            for j in range(relu_der_f.shape[0]):
                array_relu_der.append(relu_der_f[j])
            jacobian = np.diag(array_relu_der) @ jacobian
        else:
            x_layer = y_layer

        if useVars:
            m.update()

    return x_layer, jacobian[0]


def verify_initCond(wb_1_list, wb_2_list, wb_length):
    cex_init = []
    res_init = []
    m = gp.Model()
    m.setParam('Outputflag', 0)

    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
    for i in range(DIM):
        m.addConstr(x0[i] >= init_min[i])
        m.addConstr(x0[i] <= init_max[i])

    # if init_shape == 2:
    #     p0 = (init_min[0] + init_max[0]) / 2.0
    #     p1 = (init_min[1] + init_max[1]) / 2.0
    #     p2 = (init_min[2] + init_max[2]) / 2.0
    #     r = p0 - init_min[0]
    #     m.addConstr((x0[0] - p0) * (x0[0] - p0) + (x0[1] - p1) * (x0[1] - p1) + (x0[2] - p2) * (x0[2] - p2) <= r * r)

    # f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
    f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)

    # xt = np.random.randn(DIM)
    xt = [0, 0, 0, 0, 0, 0]
    result_pre = 0

    while True:
        # print("================")
        f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
        # f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)

        # print("f1_xt:", f1_xt)
        # print("df1dx_xt[0]:", df1dx_xt)
        # print("f2_x0:", f2_x0)

        m.setObjective(f1_xt
                       + df1dx_xt[0] * (x0[0] - xt[0]) + df1dx_xt[1] * (x0[1] - xt[1]) + df1dx_xt[2] * (x0[2] - xt[2]) + df1dx_xt[3] * (x0[3] - xt[3]) + df1dx_xt[4] * (x0[4] - xt[4]) + df1dx_xt[5] * (x0[5] - xt[5])
                       - f2_x0,
                       GRB.MAXIMIZE)
        m.optimize()
        # m.computeIIS()
        # m.write("model1.ilp")
        # print("m.status:", m.status)

        # print("x0.x:", x0.x)
        # print("m.objVal:", m.objVal)
        if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
            break
        xt = x0.x
        # print("xt[0]:", xt[0])
        # print("xt[1]:", xt[1])
        result_pre = m.objVal

    # print("maximize init value:", m.objVal)  # -0.07222615058680137
    if m.objVal > 0:
        cex_init.append(np.array(x0.x))
        res_init.append(m.objVal)
    return cex_init, res_init


def verify_unsafeCond(wb_1_list, wb_2_list, wb_length):
    cex_unsafe = []
    res_unsafe = []
    m = gp.Model()
    m.setParam('Outputflag', 0)

    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
    for i in range(DIM):
        m.addConstr(x0[i] >= unsafe_min[i])
        m.addConstr(x0[i] <= unsafe_max[i])

    # if unsafe_shape == 2:
    #     p0 = (unsafe_min[0] + unsafe_max[0]) / 2.0
    #     p1 = (unsafe_min[1] + unsafe_max[1]) / 2.0
    #     p2 = (unsafe_min[2] + unsafe_max[2]) / 2.0
    #     r = p0 - unsafe_min[0]
    #     m.addConstr((x0[0] - p0) * (x0[0] - p0) + (x0[1] - p1) * (x0[1] - p1) + (x0[2] - p2) * (x0[2] - p2) <= r * r)

    f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
    # f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)

    # xt = np.random.randn(DIM)
    xt = [0, 0, 0, 0, 0, 0]
    result_pre = 0

    while True:
        # print("================")
        # f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
        f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)

        # print("f1_xt:", f1_xt)
        # print("df1dx_xt[0]:", df1dx_xt)
        # print("f2_x0:", f2_x0)

        m.setObjective(f1_x0
                       - f2_xt
                       - df2dx_xt[0] * (x0[0] - xt[0]) - df2dx_xt[1] * (x0[1] - xt[1]) - df2dx_xt[2] * (x0[2] - xt[2]) - df2dx_xt[3] * (x0[3] - xt[3]) - df2dx_xt[4] * (x0[4] - xt[4]) - df2dx_xt[5] * (x0[5] - xt[5]),
                       GRB.MINIMIZE)
        m.optimize()
        # print("x0.x:", x0.x)
        # print("m.objVal:", m.objVal)
        if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
            break
        xt = x0.x
        # print("xt[0]:", xt[0])
        # print("xt[1]:", xt[1])
        result_pre = m.objVal

    # print("minimize unsafe value:", m.objVal)  # 0.36436068000596666
    if m.objVal < 0:
        cex_unsafe.append(np.array(x0.x))
        res_unsafe.append(m.objVal)
    return cex_unsafe, res_unsafe


def verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len):
    cex_domain = []
    res_domain = []

    data_dir = "./data/"
    pieces = int(10 / x_len)
    datas_xy = np.loadtxt(data_dir + f"0_10_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 4])
    datas_xxx = np.loadtxt(data_dir + f"03_0_10_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 3])

    x0_min = domain_min[0]
    x0_max = domain_max[0]
    x1_min = domain_min[1]
    x1_max = domain_max[1]
    x2_min = domain_min[2]
    x2_max = domain_max[2]
    x3_min = domain_min[3]
    x3_max = domain_max[3]
    x4_min = domain_min[4]
    x4_max = domain_max[4]
    x5_min = domain_min[5]
    x5_max = domain_max[5]

    cex_idx_list = []
    if len(cex_idx_list) == 0:
        for x0_idx in range(pieces):
            for x1_idx in range(pieces):
                for x2_idx in range(pieces):
                    for x3_idx in range(pieces):
                        for x4_idx in range(pieces):
                            for x5_idx in range(pieces):
                                cex_idx = [x0_idx, x1_idx, x2_idx, x3_idx, x4_idx, x5_idx]
                                cex_idx_list.append(cex_idx)

    for cex_idx in cex_idx_list:
        x0_idx, x1_idx, x2_idx, x3_idx, x4_idx, x5_idx = cex_idx

        k_0, b_l_0, b_0 = datas_xxx[x0_idx]
        k_1, b_l_1, b_1 = datas_xxx[x1_idx]
        k_3, b_l_3, b_3 = datas_xxx[x3_idx]
        k_4, b_l_4, b_4 = datas_xxx[x4_idx]

        w0_02, w1_02, b_l_02, b_02 = datas_xy[x0_idx * piece_len + x2_idx]
        w0_03, w1_03, b_l_03, b_03 = datas_xy[x0_idx * piece_len + x3_idx]
        w0_23, w1_23, b_l_23, b_23 = datas_xy[x2_idx * piece_len + x3_idx]
        w0_25, w1_25, b_l_25, b_25 = datas_xy[x2_idx * piece_len + x5_idx]
        w0_35, w1_35, b_l_35, b_35 = datas_xy[x3_idx * piece_len + x5_idx]

        x0_l = x0_idx * piece_len + x0_min
        x0_r = (x0_idx + 1) * piece_len + x0_min
        x1_l = x1_idx * piece_len + x1_min
        x1_r = (x1_idx + 1) * piece_len + x1_min
        x2_l = x2_idx * piece_len + x2_min
        x2_r = (x2_idx + 1) * piece_len + x2_min
        x3_l = x3_idx * piece_len + x3_min
        x3_r = (x3_idx + 1) * piece_len + x3_min
        x4_l = x4_idx * piece_len + x4_min
        x4_r = (x4_idx + 1) * piece_len + x4_min
        x5_l = x5_idx * piece_len + x5_min
        x5_r = (x5_idx + 1) * piece_len + x5_min

        m = gp.Model()
        m.setParam('Outputflag', 0)
        # m.setParam('InfUnbdInfo', 1)

        x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
        x0[0].lb = x0_l
        x0[0].ub = x0_r
        x0[1].lb = x1_l
        x0[1].ub = x1_r
        x0[2].lb = x2_l
        x0[2].ub = x2_r
        x0[3].lb = x3_l
        x0[3].ub = x3_r
        x0[4].lb = x4_l
        x0[4].ub = x4_r
        x0[5].lb = x5_l
        x0[5].ub = x5_r

        x1 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x1', lb=-np.inf, ub=np.inf)  # x1 = F(kx+b; x0)
        x1[0].lb = x0_min
        x1[0].ub = x0_max
        x1[1].lb = x1_min
        x1[1].ub = x1_max
        x1[2].lb = x2_min
        x1[2].ub = x2_max
        x1[3].lb = x3_min
        x1[3].ub = x3_max
        x1[4].lb = x4_min
        x1[4].ub = x4_max
        x1[5].lb = x5_min
        x1[5].ub = x5_max

        var_b_0 = m.addVar(lb=b_l_0, ub=b_0 + b_l_0)
        var_b_1 = m.addVar(lb=b_l_1, ub=b_1 + b_l_1)
        var_b_3 = m.addVar(lb=b_l_3, ub=b_3 + b_l_3)
        var_b_4 = m.addVar(lb=b_l_4, ub=b_4 + b_l_4)
        var_b_02 = m.addVar(lb=b_l_02, ub=b_02 + b_l_02)
        var_b_03 = m.addVar(lb=b_l_03, ub=b_03 + b_l_03)
        var_b_23 = m.addVar(lb=b_l_23, ub=b_23 + b_l_23)
        var_b_25 = m.addVar(lb=b_l_25, ub=b_25 + b_l_25)
        var_b_35 = m.addVar(lb=b_l_35, ub=b_35 + b_l_35)

        m.addConstr(x1[0] == x0[0] + time_step * (- (k_0 * x0[0] + var_b_0) + 4 * (k_1 * x0[1] + var_b_1) - 6 * (w0_23 * x0[2] + w1_23 * x0[3] + var_b_23)))
        m.addConstr(x1[1] == x0[1] + time_step * (- x0[0] - x0[1] + (k_4 * x0[4] + var_b_4)))
        m.addConstr(x1[2] == x0[2] + time_step * ((w0_03 * x0[0] + w1_03 * x0[3] + var_b_03) - x0[2] + (w0_35 * x0[3] + w1_35 * x0[5] + var_b_35)))
        m.addConstr(x1[3] == x0[3] + time_step * ((w0_02 * x0[0] + w1_02 * x0[2] + var_b_02) + (w0_25 * x0[2] + w1_25 * x0[5] + var_b_25) - (k_3 * x0[3] + var_b_3)))
        m.addConstr(x1[4] == x0[4] + time_step * (-2 * (k_1 * x0[1] + var_b_1) - x0[4] + x0[5]))
        m.addConstr(x1[5] == x0[5] + time_step * (-3 * (w0_23 * x0[2] + w1_23 * x0[3] + var_b_23) - (k_4 * x0[4] + var_b_4) - x0[5]))

        f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
        # f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)
        # f1_x1, df1dx_x1 = barrier_f(m, x1, wb_1_list, wb_length, True)
        f2_x1, df2dx_x1 = barrier_f(m, x1, wb_2_list, wb_length, True)
        # f1_x2, df1dx_x2 = barrier_f(m, x2, wb_1_list, wb_length, True)
        # f2_x2, df2dx_x2 = barrier_f(m, x2, wb_2_list, wb_length, True)

        # xt = np.random.randn(DIM)
        xt = [0, 0, 0, 0, 0, 0]
        result_pre = 0

        while True:
            # _1 means use piecewise function  # x0_next_1 == x1
            # x0_next_1 = [x0[0] + time_step * (- x0[0] + (w0 * x0[0] + w1 * x0[1] + var_b)),
            #              x0[1] + time_step * (- x0[1])]

            # _2 means use origin function
            xt_next_2 = [xt[0] + time_step * (- xt[0] * xt[0] * xt[0] + 4 * xt[1] * xt[1] * xt[1] - 6 * xt[2] * xt[3]),
                       xt[1] + time_step * (- xt[0] - xt[1] + xt[4] * xt[4] * xt[4]),
                       xt[2] + time_step * (xt[0] * xt[3] - xt[2] + xt[3] * xt[5]),
                       xt[3] + time_step * (xt[0] * xt[2] + xt[2] * xt[5] - xt[3] * xt[3] * xt[3]),
                       xt[4] + time_step * (-2 * xt[1] * xt[1] * xt[1] - xt[4] + xt[5]),
                       xt[5] + time_step * (-3 * xt[2] * xt[3] - xt[4] * xt[4] * xt[4] - xt[5])]
            # print("xt_next_2:", xt_next_2)

            # f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
            f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)
            # print("f2_xt:", f2_xt)
            # print("df2dx_xt:", df2dx_xt)

            f1_xt_next, df1dx_xt_next = barrier_f(m, xt_next_2, wb_1_list, wb_length, False)
            # print("f1_xt_next:", f1_xt_next)
            # print("df1dx_xt_next:", df1dx_xt_next)
            # f2_xt_next, df2dx_xt_next = barrier_f(m, xt_next_2, wb_2_list, wb_length, False)

            m.setObjective(f1_xt_next
                           + df1dx_xt_next[0] * (x1[0] - xt_next_2[0]) + df1dx_xt_next[1] * (x1[1] - xt_next_2[1]) + df1dx_xt_next[2] * (x1[2] - xt_next_2[2]) + df1dx_xt_next[3] * (x1[3] - xt_next_2[3]) + df1dx_xt_next[4] * (x1[4] - xt_next_2[4]) + df1dx_xt_next[5] * (x1[5] - xt_next_2[5])
                           - f2_x1
                           - (1 - alpha) * f1_x0
                           + (1 - alpha) * (f2_xt + df2dx_xt[0] * (x0[0] - xt[0]) + df2dx_xt[1] * (x0[1] - xt[1]) + df2dx_xt[2] * (x0[2] - xt[2]) + df2dx_xt[3] * (x0[3] - xt[3]) + df2dx_xt[4] * (x0[4] - xt[4]) + df2dx_xt[5] * (x0[5] - xt[5])),
                           GRB.MAXIMIZE)
            m.optimize()

            if m.status == 2:
                if m.objVal < result_pre:
                    X = xt
                    Val = result_pre
                    break
                if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
                    X = x0.x
                    Val = m.objVal
                    break
                xt = x0.x
                # print("x0.x:", x0.x)
                # print("m.objVal:", m.objVal)
                result_pre = m.objVal
            else:
                break
        # print(f"{cex_idx} is solvered!")
        if m.status == 2:
            # print("maximize domain value:", m.objVal)
            if Val > 0:
                cex_domain.append(np.array(X))
                res_domain.append(Val)

    return cex_domain, res_domain


def verif(model, x_len):
    wb_1_list = []
    wb_2_list = []
    params = model.state_dict()
    i = 0
    for k in params:
        # 'hidden_layers1.0.weight' 'hidden_layers1.0.bias' 'hidden_layers2.0.weight' 'hidden_layers2.0.bias'
        if 'layers' in k:
            continue
        if i % 4 < 2:
            wb_1_list.append(params[k].detach().numpy())
        else:
            wb_2_list.append(params[k].detach().numpy())
        i += 1

    wb_length = int(len(wb_1_list) / 2)

    cex_init, res_init = verify_initCond(wb_1_list, wb_2_list, wb_length)
    cex_unsafe, res_unsafe = verify_unsafeCond(wb_1_list, wb_2_list, wb_length)
    cex_domain, res_domain = verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len)
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../model/cdinn_zhrc1_1_20_0_0.1_1_train.pt'), strict=True)
    verif(model, piece_len)
